﻿using System;
using System.Diagnostics;
using System.Reflection;
using System.Security.Cryptography;

namespace PingCastle.misc
{
    // note: the ROCA tester need .net 4 to invoke BigInteger.
    // this code is compatile with .net2 when loaded under .net 4 context
    // default with the .config file (.net4 then .net2)
    public class ROCAVulnerabilityTester
    {
        static bool loaded = false;
        static bool loadedWithSuccess = false;
        static Type BigIntegerType;
        static ConstructorInfo ctorBigIntegerByte;
        static ConstructorInfo ctorBigIntegerInt;
        static MethodInfo op_Modulus, op_LeftShift, op_BitwiseAnd, op_Equality, op_Explicit;

        static void Load()
        {
            try
            {
                loaded = true;
                var signature = "System.Numerics";
                Trace.WriteLine("Loading " + signature);
#pragma warning disable 618
                var ass = Assembly.LoadWithPartialName(signature);
#pragma warning restore 618
                Trace.WriteLine("Get type BigInteger");
                BigIntegerType = ass.GetType("System.Numerics.BigInteger");

                Trace.WriteLine("Loading markers");
                for (int i = 0; i < markersString.Length; i++)
                {
                    markers[i] = BigIntegerType.GetMethod("Parse", new[] { typeof(string) }).Invoke(null, new object[] { markersString[i] });
                }

                Trace.WriteLine("Loading ctorBigIntegerByte");
                ctorBigIntegerByte = BigIntegerType.GetConstructor(new[] { typeof(byte[]) });
                Trace.WriteLine("Loading ctorBigIntegerInt");
                ctorBigIntegerInt = BigIntegerType.GetConstructor(new[] { typeof(int) });
                Trace.WriteLine("Loading op_Modulus");
                op_Modulus = BigIntegerType.GetMethod("op_Modulus");
                Trace.WriteLine("Loading op_BitwiseAnd");
                op_BitwiseAnd = BigIntegerType.GetMethod("op_BitwiseAnd");
                Trace.WriteLine("Loading op_LeftShift");
                op_LeftShift = BigIntegerType.GetMethod("op_LeftShift");
                Trace.WriteLine("Loading op_Equality");
                op_Equality = BigIntegerType.GetMethod("op_Equality", new Type[] { BigIntegerType, typeof(int) });
                Trace.WriteLine("Loading op_Explicit");
                foreach (MethodInfo method in BigIntegerType.GetMethods())
                {
                    if (method.Name == "op_Explicit")
                    {
                        ParameterInfo[] paramSet = method.GetParameters();
                        if ((paramSet != null) && (paramSet.Length == 1))
                        {
                            if (paramSet[0].ParameterType == BigIntegerType)
                            {
                                if (method.ReturnType == typeof(int))
                                {
                                    op_Explicit = method;
                                    break;
                                }
                            }
                        }
                    }
                }
                if (ctorBigIntegerByte == null || ctorBigIntegerInt == null || op_Modulus == null || op_LeftShift == null || op_BitwiseAnd == null || op_Equality == null || op_Explicit == null)
                    throw new NotImplementedException();
                loadedWithSuccess = true;
            }
            catch (Exception ex)
            {
                Trace.WriteLine("Unable to load BigInt");
                Trace.WriteLine(ex.Message);
                Trace.WriteLine(ex.StackTrace);
            }
        }

        // inspired from https://github.com/crocs-muni/roca/
        public static bool IsVulnerable(RSAParameters rsaKey)
        {
            if (!loaded)
            {
                Load();
            }
            if (!loadedWithSuccess)
                return false;
            //modulus is in big endian (littre endian in raw cryptoapi)
            // biginteger expects it in little endian
            var mod = new byte[rsaKey.Modulus.Length + 1];
            for (int i = 0; i < rsaKey.Modulus.Length; i++)
            {
                mod[i] = rsaKey.Modulus[rsaKey.Modulus.Length - 1 - i];
            }

            object modulus = ctorBigIntegerByte.Invoke(new object[] { mod });

            for (int i = 0; i < prims.Length; i++)
            {
                // rem = modulus % prims[i]
                var rem = op_Modulus.Invoke(null, new object[] { modulus, ctorBigIntegerInt.Invoke(new object[] { prims[i] }) });
                // t = (int) rem
                int t = (int)op_Explicit.Invoke(null, new object[] { rem });
                // shift = 1
                var shift = ctorBigIntegerInt.Invoke(new object[] { 1 });
                // shift <<= t
                shift = op_LeftShift.Invoke(null, new object[] { shift, t });
                // andresult = shift & markers[i]
                var andresult = op_BitwiseAnd.Invoke(null, new object[] { shift, markers[i] });
                // compResult = (shift & markers[i]) == 0
                var compResult = op_Equality.Invoke(null, new object[] { andresult, 0 });
                if ((bool)compResult)
                {
                    return false;
                }

            }

            return true;
        }

        private static string[] markersString = new string[]{
            "6",
            "30",
            "126",
            "1026",
            "5658",
            "107286",
            "199410",
            "8388606",
            "536870910",
            "2147483646",
            "67109890",
            "2199023255550",
            "8796093022206",
            "140737488355326",
            "5310023542746834",
            "576460752303423486",
            "1455791217086302986",
            "147573952589676412926",
            "20052041432995567486",
            "6041388139249378920330",
            "207530445072488465666",
            "9671406556917033397649406",
            "618970019642690137449562110",
            "79228162521181866724264247298",
            "2535301200456458802993406410750",
            "1760368345969468176824550810518",
            "50079290986288516948354744811034",
            "473022961816146413042658758988474",
            "10384593717069655257060992658440190",
            "144390480366845522447407333004847678774",
            "2722258935367507707706996859454145691646",
            "174224571863520493293247799005065324265470",
            "696898287454081973172991196020261297061886",
            "713623846352979940529142984724747568191373310",
            "1800793591454480341970779146165214289059119882",
            "126304807362733370595828809000324029340048915994",
            "11692013098647223345629478661730264157247460343806",
            "187072209578355573530071658587684226515959365500926"
        };

        private static object[] markers = new object[markersString.Length];

        private static int[] prims = new int[] { 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151, 157, 163, 167 };

    }
}
